#include "head.cuh"

__global__ void init_M(float *d_M, const float *d_UIview, const int nu, const int ni, const int r){
	int bid = blockIdx.x;
	int tid = threadIdx.x;
	int blkszx = blockDim.x;
	__shared__ float shmm[1024];
	shmm[tid] = 0.0f;
	__syncthreads();
	float temp = 0.f;
	for (int i = 0; i*blkszx + tid < nu; ++i){
		temp += d_UIview[i*blkszx + tid + bid*nu];
	}
	shmm[tid] = temp;
	__syncthreads();
	if (tid == 0){
		temp = 0.0f;
		for (int i = 0; i < 1024; i++){
			temp += shmm[i];
		}
		d_M[bid*r] = temp / nu;
	}
}

__global__ void get_d_Ai_d_Vi_FOR_Mi(float *d_Ai, float *d_Vi, const int r,
	const float *d_UIview, const int nu, const int ni, 
	const float *d_U,
	const int i,
	const float lambda){
	int row = threadIdx.y + blockDim.y*blockIdx.y;
	int col = threadIdx.x + blockDim.x*blockIdx.x;
	if (row < r&&col < r){
		float temp = 0.0f;
		int nui = 0;
		float temp_v = 0.0f;
		for (int j = 0; j < nu; j++){
			if (d_UIview[j+nu*i] > 0){
				nui++;
				temp += d_U[row + r*j] * d_U[col + r*j];
				temp_v += d_U[row + r*j] * d_UIview[j+nu*i];
			}
		}
		if (row == col){
			temp += lambda*nui;
		}
		d_Ai[row + col*r] = temp;
		if (col == 0){
			d_Vi[row] = temp_v;
		}
	}

}

__global__ void get_d_Ai_d_Vi_FOR_Ui(float *d_Ai, float *d_Vi, const int r,
	const float *d_UIview, const int nu, const int ni, 
	const float *d_M, 
	const int i, 
	const float lambda){
	int row = threadIdx.y + blockDim.y*blockIdx.y;
	int col = threadIdx.x + blockDim.x*blockIdx.x;
	if (row < r&&col < r){
		float temp = 0.0f;
		int nui = 0;
		float temp_v = 0.0f;
		for (int j = 0; j < ni; j++){
			if (d_UIview[i + j*nu] > 0){
				nui++;
				temp += d_M[row + r*j] * d_M[col + r*j];
				if (col==0)
					temp_v += d_M[row + r*j] * (2 * ((d_UIview[i + j*nu] - 1) / (5 - 1)) - 1)*r;
			}
		}
		if (row == col){
			temp += lambda*nui;
		}
		d_Ai[row + col*r] = temp;
		if (col == 0){
			d_Vi[row] = temp_v;
		}
	}
}

void solve_linear_system(float *d_A, float *d_B, const int m){
	cusolverDnHandle_t cusolverH;
	cusolverDnCreate(&cusolverH);

	int lwork;
	CHECK_SOLVER(cusolverDnSgeqrf_bufferSize(cusolverH, m, m, d_A, m, &lwork));
	float *d_work;
	CHECK_CUDA(cudaMalloc((void**)&d_work, sizeof(float)*lwork));
	float *d_tau;
	CHECK_CUDA(cudaMalloc(&d_tau, sizeof(float) * m));
	int *devInfo;
	CHECK_CUDA(cudaMalloc(&devInfo, sizeof(int)));
	CHECK_SOLVER(cusolverDnSgeqrf(cusolverH, m, m, d_A, m, d_tau, d_work, lwork, devInfo));
	int info_gpu;
	CHECK_CUDA(cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost));
	CHECK_PARAM(info_gpu == 0, "1 error info gpu.");
	CHECK_SOLVER(cusolverDnSormqr(cusolverH, CUBLAS_SIDE_LEFT, CUBLAS_OP_T, m, 1, m, d_A, m, d_tau, d_B, m,
		d_work, lwork, devInfo));
	CHECK_CUDA(cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost));
	CHECK_PARAM(info_gpu == 0, "2 error info gpu.");
	cublasHandle_t cublashandle;
	CHECK_BLAS(cublasCreate(&cublashandle));
	float alpha = 1.0f;
	CHECK_BLAS(cublasStrsm(cublashandle,
		CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER,
		CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT,
		m, 1,
		&alpha,
		d_A, m,
		d_B, m));
	if (d_tau) cudaFree(d_tau);
	if (devInfo) cudaFree(devInfo);
	if (d_work) cudaFree(d_work);
	if (cublashandle) CHECK_BLAS(cublasDestroy(cublashandle));
	if (cusolverH) CHECK_SOLVER(cusolverDnDestroy(cusolverH));
}

int main(void){
	cudaError_t cudaStatus = cudaSetDevice(0);
	if (cudaStatus != cudaSuccess) {
		fprintf(stderr, 
			"cudaSetDevice failed! Do you have a CUDA-capable GPU installed?");
		system("pause"); exit(1);
	}
	printf("Now: %s\n", UIviewFileName);
	clock_t start_of_main = clock();


	float *UIview = NULL;
	float split_rate = 0.8f;
	float *UIview_test = NULL;
	std::stringstream ss;

	UIview = (float*)malloc(nUsers*nItems*sizeof(float));
	CHECK_PARAM(UIview != NULL, "allocate memory space for UIview failed.");
	UIview_test = (float*)malloc(nUsers*nItems*sizeof(float));
	CHECK_PARAM(UIview_test != NULL, "allocate memory space for UIview_test failed.");
	CHECK_PARAM(read_data_2_matrix(UIview, UIview_test, split_rate), "read data set failed.");
	ss.str("");
	ss << nUsers << "_" << nItems << "_" << "UIview_test.txt";
	write_test_data(UIview_test, nUsers, nItems, ss.str().c_str());
	free(UIview_test); UIview_test = NULL;
	
	clock_t end_of_read_dataset = clock();
	printf("read data set done. cost time:%.2f s\n", 
		(double)(end_of_read_dataset - start_of_main) / CLOCKS_PER_SEC);
	
	float *d_UIview;
	CHECK_CUDA(cudaMalloc(&d_UIview, nUsers*nItems*sizeof(float)));
	CHECK_CUDA(cudaMemcpy(d_UIview, UIview, nUsers*nItems*sizeof(float), cudaMemcpyHostToDevice));

	int x[5] = { 8, 16, 32, 64, 128 };
	for (int round = 0; round < 5; round++){
		clock_t start_of_round = clock();
		printf("===================================\nround:%d %d bits\n", round, x[round]);


		int r = x[round];
		unsigned long long seed = clock() % 1000;
		float mean_of_uniform = 0.0f;
		float standard_deviation = 0.1f;
		float alpha = 1.0f;
		float beta = 0.0f;
		float lambda = 0.065f;

		curandGenerator_t gen;
		CHECK_CURAND(curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT));
		CHECK_CURAND(curandSetPseudoRandomGeneratorSeed(gen, seed));
		float *d_M;
		CHECK_CUDA(cudaMalloc(&d_M, r*nItems*sizeof(float)));
		CHECK_CURAND(curandGenerateNormal(gen, d_M, r*nItems, mean_of_uniform, standard_deviation));
		CHECK_CURAND(curandDestroyGenerator(gen));
		init_M << < nItems, 1024 >> >(d_M, d_UIview, nUsers, nItems, r);
		cudaDeviceSynchronize();
		CHECK_CUDA(cudaGetLastError());
		clock_t init_M_done = clock();
		printf("init M done. cost %.2f\n", (double)(init_M_done - start_of_round) / CLOCKS_PER_SEC);
		float *d_U;
		CHECK_CUDA(cudaMalloc(&d_U, r*nUsers*sizeof(float)));
		float *d_Ai, *d_Vi;
		CHECK_CUDA(cudaMalloc(&d_Ai, r*r*sizeof(float)));
		CHECK_CUDA(cudaMalloc(&d_Vi, r*sizeof(float)));

		cublasHandle_t cublashandle;
		CHECK_BLAS(cublasCreate(&cublashandle));//TODO something wrong here
		bool converge = false;
		int iter = 0;
		while (!converge){
			clock_t start_iter = clock();
			printf("iter:%d ",iter);
			for (int i = 0; i < nUsers; i++){
				dim3 szBlock(8, 8, 1);
				dim3 szGrid((r + 8 - 1) / 8, (r + 8 - 1) / 8, 1);
				get_d_Ai_d_Vi_FOR_Ui << <szGrid, szBlock >> >(d_Ai, d_Vi, r, d_UIview, nUsers, nItems, d_M, i, lambda);
				solve_linear_system(d_Ai, d_Vi, r);
				CHECK_CUDA(cudaMemcpy(d_U + r*i, d_Vi, r*sizeof(float), cudaMemcpyDeviceToDevice));
			}
			
			clock_t end_of_users = clock();
			printf("users done... cost %.2f ", (double)(end_of_users - start_iter) / CLOCKS_PER_SEC);
			for (int i = 0; i < nItems; i++){
				dim3 szBlock(8, 8, 1);
				dim3 szGrid((r + 8 - 1) / 8, (r + 8 - 1) / 8, 1);
				get_d_Ai_d_Vi_FOR_Mi << <szGrid, szBlock >> >(d_Ai, d_Vi, r, d_UIview, nUsers, nItems, d_U, i, lambda);
				solve_linear_system(d_Ai, d_Vi, r);
				CHECK_CUDA(cudaMemcpy(d_M + r*i, d_Vi, r*sizeof(float), cudaMemcpyDeviceToDevice));
			}
			clock_t end_of_iter = clock();
			printf("items done... cost %.2f  ",
				(double)(end_of_iter - end_of_users) / CLOCKS_PER_SEC);

			if (iter >= 50){
				converge = true;
			}
			iter++;
			printf("iter cost %.2f\n",
				(double)(end_of_iter - start_iter) / CLOCKS_PER_SEC);
		}

		CHECK_BLAS(cublasDestroy(cublashandle));

		CHECK_CUDA(cudaFree(d_Ai));
		CHECK_CUDA(cudaFree(d_Vi));

		float *B = (float *)malloc(r*nUsers*sizeof(float));
		CHECK_PARAM(B!=NULL,"allocate space for B failed.");
		CHECK_CUDA(cudaMemcpy(B, d_U, r*nUsers*sizeof(float), cudaMemcpyDeviceToHost));
		CHECK_CUDA(cudaFree(d_U)); d_U = NULL;

		ss.str("");
		ss << nUsers << "_" << nItems << "_" << r << "_" << "B_code.txt";
		write_into_file(B, r, nUsers, ss.str().c_str());
		free(B); B = NULL;

		float *D = (float *)malloc(r*nItems*sizeof(float));
		CHECK_PARAM(D != NULL, "allocate space for D failed.");
		CHECK_CUDA(cudaMemcpy(D, d_M, r*nItems*sizeof(float), cudaMemcpyDeviceToHost));
		CHECK_CUDA(cudaFree(d_M)); d_M = NULL;

		ss.str("");
		ss << nUsers << "_" << nItems << "_" << r << "_" << "D_code.txt";
		write_into_file(D, r, nItems, ss.str().c_str());
		free(D); D=NULL;
		
		clock_t end_of_round = clock();
		printf("BCCF cost time %.2f\n", (double)(end_of_round - start_of_round) / CLOCKS_PER_SEC);

	}
	clock_t end_of_main = clock();
	printf("total cost time %.2f\n", (double)(end_of_main - start_of_main) / CLOCKS_PER_SEC);
	
	CHECK_CUDA(cudaFree(d_UIview));
	

	CHECK_CUDA(cudaDeviceReset());
	system("pause");
	return EXIT_SUCCESS; 
}